#!/usr/bin/python

"""
Author: Jimmy Saw
Updated: 08-04-2012
This program draws genome comparison plot between two given genomes. It does similar things to
Artemis-based genome comparison at DNA level but instead of using DNA sequence, the program
uses reciprocal orthologous BLAST hits to chart where the orthologs between two given genomes
are located. This helps to visualize if any syntenous regions for a certain class of proteins
are present and also helps to se genome rearrangements.

Usage: 
Examples: 
Run in this folder: /host/Users/JS/UH-work/gloeobacter/final_work/comparisons/

dissertation_ReciprocalBestHitPlotWithPtt.py 58123 58319 NC_008319.ptt NC_007516.ptt orthologs/orthomcl/CYANO.orthologs.txt orthologs/orthomcl/cogs.t.list

Orthologous pair file should contain filtered orthologous
groups that should look like this:
CYANO15211: 57767|Npun_F1941 58043|Ava_0214 57803|alr2398 GKIL|GKIL_3674 59025|PCC7424_4452 58011|glr3849
CYANO15212: 57767|Npun_R1991 GKIL|GKIL_3611 58167|AM1_4177 57767|Npun_F2216 59435|Cyan7425_3596 58011|glr1283
CYANO15213: 58011|glr0952 59435|Cyan7425_0962 GKIL|GKIL_0709 57803|alr7163 58043|Ava_B0061 58167|AM1_2965
CYANO15215: 57767|Npun_R2251 GKIL|GKIL_4327 58011|glr1730 57767|Npun_R5324 58043|Ava_2173 57925|Tery_0799
"""

import sys
import re
import matplotlib.pyplot as plt
import pylab
import matplotlib
from matplotlib import mpl
from matplotlib.patches import Rectangle
from matplotlib.transforms import Bbox

g1id = sys.argv[1]

g2id = sys.argv[2]

##Regex
g1patt = re.compile('58123\|(\w+)')
g2patt = re.compile('58319\|(\w+)') # 2nd genome to compare with. can change to see different plots

cogcat = re.compile('\[(.*)\]\t(\w+)\t.*')
pttcog = re.compile('(COG\d{4})(\w).*')

#hex colors

cogdict = {
    'J' : '#2B60DE', 'A' : '#F6358A', 'K' : '#B048B5', 'L' : '#8E35EF',
    'B' : '#D16587', 'D' : '#C38EC7', 'Y' : '#52F3FF', 'V' : '#3EA99F',
    'T' : '#254117', 'M' : '#41A317', 'N' : '#00FF00', 'Z' : '#FFFF00',
    'W' : '#FDD017', 'U' : '#F88017', 'O' : '#F660AB', 'C' : '#FF0000',
    'G' : '#FAAFBA', 'E' : '#7F5A58', 'F' : '#C8B560', 'H' : '#D2B9D3',
    'I' : '#C12869', 'P' : '#57E964', 'Q' : '#BCE954', 'R' : '#F87431',
    'S' : '#ADA96E', '-' : '#EDEDED'
}

##Genome 1 ptt file
g1cogs = {}
g1pttfile = open(sys.argv[3], "rU")
g1ptt = g1pttfile.readlines()
tmp1 = g1ptt[0]
genome1size = int(tmp1.split(' ')[-1].split('..')[1])
g1featdict = {}
g1orgname = g1ptt[0].split(',')[0]

for line in g1ptt[3:]:
    c = line.split('\t')
    g1ltag = c[5]
    g1gene = c[4]
    g1start = int(c[0].split('..')[0])
    g1stop = int(c[0].split('..')[1])
    g1pid = c[3]
    g1def = c[-1].strip()
    g1strand = c[1]
    g1cog = '-'
    if pttcog.match(c[7]):
        p = pttcog.match(c[7])
        g1cog = p.group(1)
    g1cogs[g1ltag] = g1cog
    g1featdict[g1ltag] = ((g1start, g1stop, g1strand, g1pid, g1def, g1cog))

##Genome 2 ptt file
g2cogs = {}
g2pttfile = open(sys.argv[4], "rU")
g2ptt = g2pttfile.readlines()
tmp2 = g2ptt[0]
genome2size = int(tmp2.split(' ')[-1].split('..')[1])
g2featdict = {}
g2orgname = g2ptt[0].split(',')[0]

for line in g2ptt[3:]:
    c = line.split('\t')
    g2ltag = c[5]
    g2gene = c[4]
    g2start = int(c[0].split('..')[0])
    g2stop = int(c[0].split('..')[1])
    g2pid = c[3]
    g2def = c[-1].strip()
    g2strand = c[1]
    g2cog = '-'
    if pttcog.match(c[7]):
        p = pttcog.match(c[7])
        g2cog = p.group(1)
    g2cogs[g2ltag] = g2cog
    g2featdict[g2ltag] = ((g2start, g2stop, g2strand, g2pid, g2def, g2cog))

##Compare genome sizes
larger_genome = 0

if genome1size > genome2size:
    larger_genome = genome1size
else:
    larger_genome = genome2size

##Pair file
pairfile = open(sys.argv[5], "rU")
pfl = pairfile.readlines()

totalg1hits = 0
totalg2hits = 0

cogcatfile = open(sys.argv[6], "rU")
cfl = cogcatfile.readlines()

cogcatdict = {}

for line in cfl:
    tmp = line.strip()
    if cogcat.match(tmp):
        pattern = cogcat.match(tmp)
        cogcatdict[pattern.group(2)] = pattern.group(1)[0]

glist = []

for line in pfl:
    t = line.split(' ')
    oname = t[0]
    g1x = 0 #x-coord for genome1
    g2x = 0 #x-coord for genome2
    g1cogcolor = "#EDEDED"
    g2cogcolor = "#EDEDED"
    for i in t[1:]: #note that each line may contain multiple hits. eg: 2 vs 3
        tmp = i.strip()
        if g1patt.match(tmp):
            pat1 = g1patt.match(tmp)
            ltag = pat1.group(1)
            if ltag in g1featdict:
                g1feat = g1featdict[ltag]
                g1start = g1featdict[ltag][0]
                g1stop = g1featdict[ltag][1]
                g1mid = g1start + ((g1stop - g1start)/2.0) #places the point to middle of gene
                g1x = g1mid
                g1desc = g1featdict[ltag][-2]
                g1cog = g1featdict[ltag][-1]
                if g1cog in cogcatdict:
                    g1cogcolor = cogdict[cogcatdict[g1cog]]
                #print '{0:10}\t{1}\t{2}\t{3}'.format(ltag, g1start, g1stop, g1desc)
                totalg1hits += 1
        if g2patt.match(tmp):
            pat2 = g2patt.match(tmp)
            ltag = pat2.group(1)
            if ltag in g2featdict:
                g2feat = g2featdict[ltag]
                g2start = g2featdict[ltag][0]
                g2stop = g2featdict[ltag][1]
                g2mid = g2start + ((g2stop - g2start)/2.0) #places the point to middle of gene
                g2x = g2mid
                g2desc = g2featdict[ltag][-2]
                g2cog = g2featdict[ltag][-1]
                if g2cog in cogcatdict:
                    g2cogcolor = cogdict[cogcatdict[g2cog]]
                #print '{0:10}\t{1}\t{2}\t{3}'.format(ltag, g1start, g1stop, g1desc)
                totalg2hits += 1
    glist.append((g1x, g2x, g1cogcolor)) #append the x-coordinates for genome 1 and 2, and cog color

##For lines representing the genomes. Needed to set upper and lower bounds above and below x-y segments
genome1x = [2, genome1size]
genome1y = [2, 2]
genome2x = [2, genome2size]
genome2y = [10, 10]

##Start plotting
#fig = plt.figure(1, figsize=(15,8)) #to be used with 211
fig = plt.figure(1, figsize=(15,4))
#ax1 = fig.add_subplot(211) #makes the subplot and squeezes the figure to half panel
ax1 = fig.add_subplot(111) #makes the full figure plot. larger.
ax1.plot(genome1x, genome1y, color='#FFFFFF', marker='|', markersize=8.0,
    mec='black', ls='-', lw=2.0)
ax1.plot(genome2x, genome2y, color='#FFFFFF', marker='|', markersize=8.0,
    mec='black', ls='-', lw=2.0)
ax1.axis([0, larger_genome, 0, 14])


for i in glist:
    #print i[0], i[1]
    x = [i[0], i[1]] #i[0] is genome 1 coord and i[1] is genome 2 coord
    y = [3, 9] #y axes
    c = i[2]
    #ax1.plot(x, y, "o-", alpha=0.2)
    #ax1.plot(x, y, "s-", alpha=0.2)
    #ax1.plot(x, y, color=c, marker="o", alpha=0.2)
    ax1.plot(x, y, color=c, marker="s", alpha=0.2)
    #ax1.plot(x, y, "#333366", alpha=0.2) #should set the color based on COG categories

ax1.annotate(g1orgname, xy=(0.7, 0.05), xycoords='axes fraction', 
    horizontalalignment='center', verticalalignment='center', fontsize=10)
ax1.annotate(g2orgname, xy=(0.7, 0.93), xycoords='axes fraction', 
    horizontalalignment='center', verticalalignment='center', fontsize=10)

#Draw legend box for COG categories
coglist = []
for k, v in cogdict.iteritems():
    coglist.append((k,v))

coglist.sort()

ccounts = len(coglist)
a = 0
xstart = 500000
#increment = 20000 #for synecoccus
increment = 40000 #for gloeobacter
while a < ccounts:
    fc = coglist[a][1]
    tx = coglist[a][0]
    #rect = Rectangle((xstart, 2.5), 20000, 0.25, facecolor=fc, alpha=0.5) #for synecococcus
    rect = Rectangle((xstart, 2.5), 40000, 0.25, facecolor=fc, alpha=0.5) #for gloeobacter
    plt.gca().add_patch(rect)
    ax1.annotate(tx, xy=(xstart+20000, 2.3), horizontalalignment='center', 
        verticalalignment='center', fontsize=8)
    a += 1
    xstart = xstart + increment
ax1.annotate('COG categories', xy=(100000, 2.3), fontsize=8)

frame1 = plt.gca()
for tick in frame1.axes.get_yticklines():
    tick.set_visible(False)
for y in frame1.axes.get_yticklabels():
    y.set_visible(False)
ax1.grid(False)
plt.show()

pairfile.close()


